import numpy as np
import os
import tensorflow as tf
import cv2
from tensorflow.python.platform import gfile
import time
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' 
os.environ["CUDA_VISIBLE_DEVICES"] = "0"  

class rcnn():
    def __init__(self,pb_file):
        config = tf.ConfigProto()
        config.gpu_options.allow_growth = True
        self.sess = tf.Session(config=config)
        with gfile.FastGFile(pb_file, 'rb') as f:
            graph_def = tf.GraphDef()
            graph_def.ParseFromString(f.read())
            self.sess.graph.as_default()
            tf.import_graph_def(graph_def, name='')
    def get_box(self, img):
        img = cv2.medianBlur(img, 7)
        image = np.expand_dims(img,axis = 0)
        boxes,scores,classes,num_detections = self.sess.run([self.sess.graph.get_tensor_by_name('detection_boxes:0'),self.sess.graph.get_tensor_by_name('detection_scores:0'),self.sess.graph.get_tensor_by_name('detection_classes:0'),self.sess.graph.get_tensor_by_name('num_detections:0')],
                                                            feed_dict={self.sess.graph.get_tensor_by_name('image_tensor:0'):image})
        box,classes_list = self.re_norm_boxes(np.squeeze(boxes),np.squeeze(scores),np.squeeze(classes),img.shape,0.4,0.2)
        bac1,bac2,bac3,bac4 = self.classes_box(box, classes_list)
        return bac1,bac2,bac3,bac4

    def classes_box (self, boxes, classes_list):
        bac1 = []
        bac2 = []
        bac3 = []
        bac4 = []

        for i in range(len(boxes)):
            box = boxes[i]
            classes = classes_list[i]
            if classes == 1.0:
                bac1.append([box[0],box[1],box[2],box[3]])
            elif classes == 2.0:
                bac2.append([box[0],box[1],box[2],box[3]])
            elif classes == 3.0:
                bac3.append([box[0],box[1],box[2],box[3]])
            elif classes == 4.0:
                bac4.append([box[0],box[1],box[2],box[3]])

        return bac1,bac2,bac3,bac4


    def re_norm_boxes(self,boxes,scores,classes,img_shape,scores_hreshold=0.5,hreshold=0.8):
        box = []
        final_box = []
        classes_list = []
        def j(dets, thresh):
            x1 = dets[:,0]
            y1 = dets[:,1]
            x2 = dets[:,2]
            y2 = dets[:,3]
            areas = (y2-y1+1) * (x2-x1+1)
            scores = dets[:,4]
            keep = []
            index = scores.argsort()[::-1]
            while index.size >0:
                i = index[0]
                keep.append(i)
                x11 = np.maximum(x1[i], x1[index[1:]])
                y11 = np.maximum(y1[i], y1[index[1:]])
                x22 = np.minimum(x2[i], x2[index[1:]])
                y22 = np.minimum(y2[i], y2[index[1:]])
                w = np.maximum(0, x22-x11+1)    
                h = np.maximum(0, y22-y11+1)    
                overlaps = w*h
                ious = overlaps / (areas[i]+areas[index[1:]] - overlaps)
                idx = np.where(ious<=thresh)[0]
                index = index[idx+1]
            return keep

        for i in range(len(boxes)):
            if scores[i] > scores_hreshold:
                ymin,xmin,ymax,xmax = boxes[i]
                box.append([int(xmin*img_shape[1]),int(ymin*img_shape[0]),int(xmax*img_shape[1]),int(ymax*img_shape[0]),scores[i]])


        _index = j(np.asarray(box),hreshold)
        for index in _index:
            final_box.append(box[index])
            classes_list.append(classes[index])

        return final_box,classes_list


    




if __name__ == '__main__':
    
    pb_dir = './jljs1.pb'
    img_dir = './jljs/'

    use_object_detect_pb = rcnn(pb_dir)
    img_list = os.listdir(img_dir)
    all_time = 0
    for i in range(len(img_list)):
        t1 = time.time()
        print(img_list[i])
        img = cv2.imread(img_dir+img_list[i])
        # img = cv2.blur(img, (5, 5))
        # img = cv2.medianBlur(img, 5)
        bac1,bac2,bac3,bac4 = use_object_detect_pb.get_box(img)
        for j in range(len(bac1)):
            box  = bac1[j]
            img = cv2.rectangle(img, (box[0], box[1]), (box[2], box[3]), (0, 0, 255), 1)
        for j in range(len(bac2)):
            box = bac2[j]
            img = cv2.rectangle(img, (box[0], box[1]), (box[2], box[3]), (0, 255, 0), 1)
        for j in range(len(bac3)):
            box = bac3[j]
            img = cv2.rectangle(img, (box[0], box[1]), (box[2], box[3]), (255, 0, 0), 1)
        for j in range(len(bac4)):
             ox = bac4[j]
             img = cv2.rectangle(img, (box[0], box[1]), (box[2], box[3]), (255, 255, 0), 1)
        cv2.imwrite('./out/'+img_list[i],img)